import math
import torch
import torch.nn as nn
# import torch.fx
from torch import Tensor
import torch.nn.functional as F
from collections import OrderedDict
from batchrenorm import BatchRenorm2d

class Mish(nn.Module):
    def __init__(self):
        super(Mish, self).__init__()

    def forward(self, x):
        return x * torch.tanh(F.softplus(x))


class BN_Conv_Mish(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation=1, groups=1, bias=False):
        super(BN_Conv_Mish, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation=dilation,
                              groups=groups, bias=bias)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        out = self.bn(self.conv(x))
        return Mish()(out)
    '''随机深度'''


class ResidualBlock(nn.Module):
    """
    basic residual block for CSP-Darknet
    """
    def __init__(self, chnls, inner_chnnls=None):
        super(ResidualBlock, self).__init__()
        if inner_chnnls is None:
            inner_chnnls = chnls
        self.conv1 = BN_Conv_Mish(chnls, inner_chnnls, 1, 1, 0)     # always use samepadding
        self.conv2 = nn.Conv2d(inner_chnnls, chnls, 3, 1, 1, bias=False)
        self.bn = nn.BatchNorm2d(chnls)

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.bn(out) + x
        return Mish()(out)

class CSPFirst(nn.Module):
    """
    First CSP Stage
    """
    def __init__(self, in_chnnls, out_chnls):
        super(CSPFirst, self).__init__()
        self.dsample = BN_Conv_Mish(in_chnnls, out_chnls, 3, 2, 1)     # same padding
        self.trans_0 = BN_Conv_Mish(out_chnls, out_chnls, 1, 1, 0)
        self.trans_1 = BN_Conv_Mish(out_chnls, out_chnls, 1, 1, 0)
        self.block = ResidualBlock(out_chnls, out_chnls//2)
        self.trans_cat = BN_Conv_Mish(2*out_chnls, out_chnls, 1, 1, 0)

    def forward(self, x):
        x = self.dsample(x)
        out_0 = self.trans_0(x)
        out_1 = self.trans_1(x)
        out_1 = self.block(out_1)
        out = torch.cat((out_0, out_1), 1)
        out = self.trans_cat(out)
        return out


class CSPStem(nn.Module):
    """
    CSP structures including downsampling
    """

    def __init__(self, in_chnls, out_chnls, num_block):
        super(CSPStem, self).__init__()
        self.dsample = BN_Conv_Mish(in_chnls, out_chnls, 3, 2, 1)
        self.trans_0 = BN_Conv_Mish(out_chnls, out_chnls//2, 1, 1, 0)
        self.trans_1 = BN_Conv_Mish(out_chnls, out_chnls//2, 1, 1, 0)
        self.blocks = nn.Sequential(*[ResidualBlock(out_chnls//2) for _ in range(num_block)])
        self.trans_cat = BN_Conv_Mish(out_chnls, out_chnls, 1, 1, 0)

    def forward(self, x):
        x = self.dsample(x)
        out_0 = self.trans_0(x)
        out_1 = self.trans_1(x)
        out_1 = self.blocks(out_1)
        out = torch.cat((out_0, out_1), 1)
        out = self.trans_cat(out)
        return out


class CSP_DarkNet(nn.Module):
    """
    CSP-DarkNet
    """
    def __init__(self, num_blocks: object, num_classes=1000) -> object:
        super(CSP_DarkNet, self).__init__()
        # chnls = [64, 128, 256, 512, 1024]
        chnls = [32,36,40,80,112]
        self.conv0 = BN_Conv_Mish(3, 32, 3, 1, 1)   # same padding
        self.neck = CSPFirst(32, chnls[0])
        self.body = nn.Sequential(
            *[CSPStem(chnls[i], chnls[i+1], num_blocks[i]) for i in range(4)])
        self.conv_out6_channel = nn.Conv2d(chnls[4], 192, kernel_size=3, padding=1, stride=1)
        self.conv_out6 = nn.MaxPool2d(kernel_size=3, padding=1, stride=2)
        self.conv_out7_channel = nn.Conv2d(192, 320, kernel_size=3, padding=1, stride=1)
        self.conv_out7 = nn.MaxPool2d(kernel_size=3, padding=1, stride=2)

    def forward(self, x):
        out = self.conv0(x)
        out = self.neck(out)
        out = self.body(out)

        return F.softmax(out)

# def csp_darknet_53(num_classes=1000):
#     return CSP_DarkNet([2, 8, 8, 4], num_classes)

class csp_darknet_53(nn.Module):
    def __init__(self,num_class=15):
        super(csp_darknet_53, self).__init__()
        model = CSP_DarkNet([2, 8, 8, 4],num_classes=num_class)
        # del model.stage7_conv
        # del model.stage7_globalAverPool
        # del model.stage7_classfier
        self.model = model
    def forward(self,x):
        # x = self.model.stage0_conv3(x)
        # out1 = self.model.stage1to6[0](x)
        # out2 = self.model.stage1to6[1](out1)
        # out3 = self.model.stage1to6[2](out2)
        # out4 = self.model.stage1to6[3](out3)
        # out5 = self.model.stage1to6[4](out4)
        # out6 = self.model.stage1to6[5](out5)
        # return out4,out5,out6
        x = self.model.conv0(x)
        x = self.model.neck(x)

        feature_maps = []
        for idx, stage in enumerate(self.model.body):
            x = stage(x)
            feature_maps.append(x)
        out6 = self.model.conv_out6_channel(feature_maps[3])
        out6 = self.model.conv_out6(out6)
        out7 = self.model.conv_out7_channel(out6)
        out7 = self.model.conv_out7(out7)
        #[64,52,52],[128,26,26],[256,13,13]
        return feature_maps[1],feature_maps[2],feature_maps[3],out6,out7